package jamirr.loader;

import java.io.*;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.FileSystem;
import java.nio.file.*;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.*;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.RecursiveAction;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;

final class ServiceScanner<S> {

    private static final String SCHEME_FILE = "file";
    private static final String SCHEME_JAR = "jar";
    private static final String SCHEME_ZIP = "zip";
    private static final String SCHEME_WSJAR = "wsjar";

    private static final String COLON = ":";


    private final ClassLoader classLoader;
    private final String serviceName;
    private final Predicate<String> lineCondition;
    private final Function<String, S> transformer;

    public ServiceScanner(ClassLoader classLoader, String serviceName, Predicate<String> lineCondition, Function<String, S> transformer) {
        this.classLoader = classLoader;
        this.serviceName = serviceName;
        this.lineCondition = lineCondition;
        this.transformer = transformer;
    }

    @SuppressWarnings({"all"})
    static void eachFile(URI uri, String path, Consumer<Path> consumer) {
        List<Closeable> toClose = new ArrayList<>();
        try {
            Path myPath = resolvePath(uri, path, toClose, ServiceScanner::loadNestedJarUri);
            if (myPath != null) {
                Path finalMyPath = myPath;
                // use this method instead of Files#walk to eliminate the Stream overhead
                Files.walkFileTree(myPath, Collections.emptySet(), 1, new FileVisitor<>() {
                    @Override
                    public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) {
                        return FileVisitResult.CONTINUE;
                    }

                    @Override
                    public FileVisitResult visitFile(Path currentPath, BasicFileAttributes attrs) throws IOException {
                        if (currentPath.equals(finalMyPath) || Files.isHidden(currentPath) || currentPath.getFileName().startsWith(".")) {
                            return FileVisitResult.CONTINUE;
                        }
                        consumer.accept(currentPath);
                        return FileVisitResult.CONTINUE;
                    }

                    @Override
                    public FileVisitResult visitFileFailed(Path file, IOException exc) {
                        return FileVisitResult.CONTINUE;
                    }

                    @Override
                    public FileVisitResult postVisitDirectory(Path dir, IOException exc) {
                        return FileVisitResult.CONTINUE;
                    }
                });
            }
        } catch (IOException e) {
            // ignore, can't do anything here and can't log because class used in compiler
        } finally {
            for (Closeable closeable : toClose) {
                try {
                    closeable.close();
                } catch (IOException ignored) {
                }
            }
        }
    }

    static Path resolvePath(URI uri,
                            String path,
                            List<Closeable> toClose,
                            IOExceptionBiFunction<List<Closeable>, String, Path> loadNestedJarUriFunction) throws IOException {
        String scheme = uri.getScheme();
        try {
            if (SCHEME_JAR.equals(scheme) || SCHEME_ZIP.equals(scheme) || SCHEME_WSJAR.equals(scheme)) {
                // try to match FileSystems.newFileSystem(URI) semantics for zipfs here.
                // Basically ignores anything after the !/ if it exists, and uses the part
                // before as the jar path to extract.
                String jarUri = uri.getRawSchemeSpecificPart();
                int sep = jarUri.lastIndexOf("!/");
                if (sep != -1) {
                    jarUri = jarUri.substring(0, sep);
                }
                if (!jarUri.startsWith(SCHEME_FILE + COLON)) {
                    // Special case WebLogic classloader
                    jarUri = jarUri.startsWith("/") ?
                            SCHEME_FILE + COLON + jarUri :
                            SCHEME_FILE + COLON + "/" + jarUri;
                }
                // now, add the !/ at the end again so that loadNestedJarUri can handle it:
                jarUri += "!/";
                return loadNestedJarUriFunction.apply(toClose, jarUri).resolve(path);
            } else if ("file".equals(scheme)) {
                return Paths.get(uri).resolve(path);
            } else {
                // graal resource: case
                return Paths.get(uri);
            }
        } catch (FileSystemNotFoundException e) {
            return null;
        }
    }

    private static Path loadNestedJarUri(List<Closeable> toClose, String jarUri) throws IOException {
        int sep = jarUri.lastIndexOf("!/");
        if (sep == -1) {
            return Paths.get(URI.create(jarUri));
        }
        Path jarPath = loadNestedJarUri(toClose, jarUri.substring(0, sep));
        if (Files.isDirectory(jarPath)) {
            // spring boot creates weird jar URLs, like 'jar:file:/xyz.jar!/BOOT-INF/classes!/abc'
            // This check makes our class loading resilient to that
            return jarPath;
        }
        FileSystem zipfs = FileSystems.newFileSystem(jarPath, (ClassLoader) null);
        toClose.addFirst(zipfs);
        return zipfs.getPath(jarUri.substring(sep + 1));
    }

    @FunctionalInterface
    public interface IOExceptionBiFunction<T, U, R> {
        R apply(T t, U u) throws IOException;
    }

    private static URI normalizeFilePath(String path, URI uri) {
        Path p = Paths.get(uri);
        if (p.endsWith(path)) {
            Path subpath = Paths.get(path);
            for (int i = 0; i < subpath.getNameCount(); i++) {
                p = p.getParent();
            }
            uri = p.toUri();
        }
        return uri;
    }

    @SuppressWarnings("java:S3398")
    private static Set<String> computeServiceTypeNames(URI uri, String path) {
        Set<String> typeNames = new HashSet<>();
        // Keep the anonymous class instead of Lambda to reduce the Lambda invocation overhead during the startup
        Consumer<Path> consumer = currentPath -> {
            if (Files.isRegularFile(currentPath)) {
                final String typeName = currentPath.getFileName().toString();
                typeNames.add(typeName);
            }
        };
        eachFile(uri, path, consumer);
        return typeNames;
    }

    @SuppressWarnings("java:S3398")
    private Set<String> computeStandardServiceTypeNames(URL url) {
        Set<String> typeNames = new HashSet<>();
        try {
            try (BufferedReader reader = new BufferedReader(new InputStreamReader(url.openStream()))) {
                while (true) {
                    String line = reader.readLine();
                    if (line == null) {
                        break;
                    }
                    if (line.isEmpty() || line.charAt(0) == '#') {
                        continue;
                    }
                    if (!lineCondition.test(line)) {
                        continue;
                    }
                    int i = line.indexOf('#');
                    if (i > -1) {
                        line = line.substring(0, i);
                    }
                    typeNames.add(line);
                }
            }
        } catch (IOException | UncheckedIOException e) {
            // ignore, can't do anything here and can't log because class used in compiler
        }
        return typeNames;
    }

    private boolean isWebSphereClassLoader() {
        return classLoader.getClass().getName().startsWith("com.ibm.ws.classloader");
    }

    private String buildResourceSearchPath() {
        String path = "META-INF/jamirr/" + serviceName;
        if (isWebSphereClassLoader()) {
            return path + "/";
        }
        return path;
    }

    private Enumeration<URL> findStandardServiceConfigs() throws IOException {
        return classLoader.getResources(ServiceLoaderImpl.META_INF_SERVICES + '/' + serviceName);
    }

    private void findMetaServiceConfigs(BiConsumer<URI, String> consumer) throws IOException, URISyntaxException {
        String path = buildResourceSearchPath();
        final Enumeration<URL> resources = classLoader.getResources(path);
        Set<URI> uniqueURIs = new LinkedHashSet<>();
        while (resources.hasMoreElements()) {
            URL url = resources.nextElement();
            final URI uri = url.toURI();
            uniqueURIs.add(uri);
        }

        for (URI uri : uniqueURIs) {
            String scheme = uri.getScheme();
            if ("file".equals(scheme)) {
                uri = normalizeFilePath(path, uri);
            }
            // on GraalVM there are spurious extra resources that end with # and then a number
            // we ignore this extra ones
            if (!("resource".equals(scheme) && uri.toString().contains("#"))) {
                consumer.accept(uri, path);
            }
        }
    }

    /**
     * Fork-join recursive service loader.
     */
    @SuppressWarnings("java:S1948")
    final class DefaultServiceCollector extends RecursiveActionValuesCollector<S> implements ServiceLoaderImpl.ServiceCollector<S> {

        private final List<RecursiveActionValuesCollector<S>> tasks = new ArrayList<>();

        @Override
        protected void compute() {
            try {
                Enumeration<URL> serviceConfigs = findStandardServiceConfigs();
                while (serviceConfigs.hasMoreElements()) {
                    URL url = serviceConfigs.nextElement();
                    UrlServicesLoader task = new UrlServicesLoader(url);
                    tasks.add(task);
                    task.fork();
                }
                findMetaServiceConfigs((uri, path) -> {
                    final MetaServicesLoader task = new MetaServicesLoader(uri, path);
                    tasks.add(task);
                    task.fork();
                });
            } catch (IOException | URISyntaxException e) {
                throw new ServiceConfigurationError("Failed to load resources for service: " + serviceName, e);
            }
        }

        @SuppressWarnings("all")
        @Override
        public void collect(Collection<S> values) {
            ForkJoinPool.commonPool().invoke(this);
            for (RecursiveActionValuesCollector<S> task : tasks) {
                task.join();
                task.collect(values);
            }
        }

        @SuppressWarnings("all")
        @Override
        public void collect(Collection<S> values, boolean allowFork) {
            if (allowFork) {
                ForkJoinPool.commonPool().invoke(this);
                for (RecursiveActionValuesCollector<S> task : tasks) {
                    task.join();
                    task.collect(values);
                }
            } else {
                try {
                    Enumeration<URL> serviceConfigs = findStandardServiceConfigs();
                    while (serviceConfigs.hasMoreElements()) {
                        URL url = serviceConfigs.nextElement();
                        for (String typeName : computeStandardServiceTypeNames(url)) {
                            S val = transformer.apply(typeName);
                            if (val != null) {
                                values.add(val);
                            }
                        }
                    }
                    findMetaServiceConfigs((uri, path) -> {
                        for (String typeName : computeServiceTypeNames(uri, path)) {
                            S val = transformer.apply(typeName);
                            if (val != null) {
                                values.add(val);
                            }
                        }
                    });
                } catch (IOException | URISyntaxException e) {
                    throw new ServiceConfigurationError("Failed to load resources for service: " + serviceName, e);
                }
            }
        }
    }

    private final class MetaServicesLoader extends RecursiveActionValuesCollector<S> {
        private final URI uri;
        private final List<ServiceInstanceLoader> tasks = new ArrayList<>();
        private final String path;

        private MetaServicesLoader(URI uri, String path) {
            this.uri = uri;
            this.path = path;
        }

        @Override
        public void collect(Collection<S> values) {
            for (ServiceInstanceLoader task : tasks) {
                task.join();
                task.collect(values);
            }
        }

        @Override
        @SuppressWarnings("java:S2095")
        protected void compute() {
            Set<String> typeNames = computeServiceTypeNames(uri, path);
            for (String typeName : typeNames) {
                ServiceInstanceLoader task = new ServiceInstanceLoader(typeName);
                tasks.add(task);
                task.fork();
            }
        }
    }

    /**
     * Reads URL, parses the file and produces sub-tasks to initialize the entry.
     */
    @SuppressWarnings("java:S1948")
    private final class UrlServicesLoader extends RecursiveActionValuesCollector<S> {

        private final URL url;
        private final List<ServiceInstanceLoader> tasks = new ArrayList<>();

        public UrlServicesLoader(URL url) {
            this.url = url;
        }

        @Override
        @SuppressWarnings({"java:S3776", "java:S135"})
        protected void compute() {
            for (String typeName : computeStandardServiceTypeNames(url)) {
                ServiceInstanceLoader task = new ServiceInstanceLoader(typeName);
                tasks.add(task);
                task.fork();
            }
        }

        @Override
        public void collect(Collection<S> values) {
            for (ServiceInstanceLoader task : tasks) {
                task.join();
                task.collect(values);
            }
        }

    }

    /**
     * Initializes and filters the entry.
     */
    @SuppressWarnings("java:S1948")
    private final class ServiceInstanceLoader extends RecursiveActionValuesCollector<S> {

        private final String className;
        private S result;
        private Throwable throwable;

        public ServiceInstanceLoader(String className) {
            this.className = className;
        }

        @Override
        protected void compute() {
            try {
                result = transformer.apply(className);
            } catch (Throwable e) {
                throwable = e;
            }
        }

        @Override
        public void collect(Collection<S> values) {
            if (throwable != null) {
                throw new ServiceLoaderImpl.ServiceLoadingException("Failed to load a service: " + throwable.getMessage(), throwable);
            }
            if (result != null && !values.contains(result)) {
                values.add(result);
            }
        }
    }

    /**
     * Abstract recursive action class.
     *
     * @param <S> The type
     */
    private abstract static class RecursiveActionValuesCollector<S> extends RecursiveAction {

        /**
         * Collects loaded values.
         *
         * @param values The values
         */
        public abstract void collect(Collection<S> values);

    }
}
